# -*- coding: utf-8 -*-
"""
LSTM model
"""
import  os
import tensorflow as tf
import  numpy as np
from tensorflow.keras import Model, layers, losses, optimizers, Sequential, models

class LSTM(Model):
    #Cell方式构建多层网络
    def __init__(self, units = 24, num_classes = 1):
        super(LSTM, self).__init__()
        #构建lstm
        self.lstm = layers.LSTM(units, dropout = 0.5)
        self.dense1 = layers.Dense(32)
        self.dense2 = layers.Dense(num_classes)
        
    def call(self, inputs, training = None):
        x = inputs
        x = self.lstm(x)
        x = self.dense1(x)
        x = tf.nn.relu(x)
        x = self.dense2(x)        
        return x
    
 
def train(units, num_classes, x_train, y_train, x_valid, y_valid, batch_size, epochs, patience, save_path):
    """
    units : lstm层神经元个数
    num_classes : 预测的label长度
    x_train : 训练集x
    y_train : 训练集y
    x_valid : 验证集x
    y_valid : 验证集y
    batch_size : 每次放入的数量
    epochs : 迭代次数
    patience : epoch没有改进的个数
    save_path : 模型保存的路径
    
    return 训练结果
    """
    #早停机制
    early_stopping = tf.keras.callbacks.EarlyStopping(monitor='val_loss',
                                                    patience = patience,
                                                    mode='min')
    #加载模型结构
    model = LSTM(units, num_classes)
    #装配模型
    model.compile(optimizer = optimizers.Adam(0.001),
                  loss = "mse",
                  metrics=['accuracy'])
    #训练
    model_result = model.fit(x=x_train, y=y_train, batch_size = batch_size, epochs = epochs, validation_data = (x_valid, y_valid), callbacks=[early_stopping])
    model.summary()
    #保存模型
    model.save_weights(save_path, save_format = "tf")  
    #训练集结果
    train_result = model.predict(x_train)
    valid_result = model.predict(x_valid)
    
    del model
    
    return model_result, train_result , valid_result
    
def predict(units, num_classes, save_path, x_test):
    """
    units : lstm层神经元个数
    num_classes : 预测的label长度
    save_path : 模型保存的路径
    x_test : 测试集
    使用权重保存模型需加载原模型结构
    
    return  预测结果
    """
    #加载原模型结构
    model = LSTM(units, num_classes)
    #装配模型
    model.compile(optimizer = optimizers.Adam(0.001),
                  loss = "mse",
                  metrics=['accuracy'])
    #加载模型权重
    model.load_weights(save_path)
    #预测
    predict_result = model.predict(x_test)
    
    del model
    
    return predict_result

if __name__ == "main":
    units = 64
    batch_size = 20
    epochs = 100
    patience = 10
    save_path = "lstm.ckpt"
